#!/usr/bin/env python3
"""
Plot training confidence values from multiple CSV files
"""

import pandas as pd
import matplotlib.pyplot as plt
import os
import numpy as np
from pathlib import Path

def plot_confidence_values():
    """Plot Value column from multiple confidence CSV files"""
    
    # Define the data directory
    data_dir = "/home/shaonian/SED/SED/script_tools/data"
    
    # Find all confidence CSV files
    csv_files = []
    for file in os.listdir(data_dir):
        if file.endswith("conficence.csv"):  # Note: keeping the original typo "conficence"
            csv_files.append(os.path.join(data_dir, file))
    
    print(f"Found {len(csv_files)} CSV files:")
    for file in csv_files:
        print(f"  - {os.path.basename(file)}")
    
    # Create figure and axis
    plt.figure(figsize=(15, 10))
    
    # Colors for different lines
    colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown', 'pink', 'gray']
    
    # Plot each file
    for i, csv_file in enumerate(csv_files):
        try:
            # Read CSV file
            df = pd.read_csv(csv_file)
            
            # Extract filename for legend
            filename = os.path.basename(csv_file)
            
            # Create a more readable label by extracting key information
            if "teacher" in filename:
                model_type = "Teacher"
            elif "student" in filename:
                model_type = "Student"
            else:
                model_type = "Unknown"
            
            if "in-domain" in filename:
                dataset_type = "In-domain"
            elif "refTrainWeak" in filename:
                dataset_type = "RefTrainWeak"
            else:
                dataset_type = "Standard"
            
            label = f"{model_type} - {dataset_type}"
            
            # Plot the values
            plt.plot(df['Step'], df['Value'], 
                    label=label, 
                    color=colors[i % len(colors)], 
                    linewidth=2,
                    alpha=0.8)
            
            print(f"Plotted {filename}: {len(df)} data points")
            
        except Exception as e:
            print(f"Error reading {csv_file}: {e}")
            continue
    
    # Customize the plot
    plt.xlabel('Training Step', fontsize=14)
    plt.ylabel('Confidence Value', fontsize=14)
    plt.title('Training Confidence Evolution Comparison', fontsize=16, fontweight='bold')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=12)
    plt.grid(True, alpha=0.3)
    
    # Set axis limits
    plt.xlim(left=0)
    plt.ylim(0, 1.0)  # Assuming confidence values are between 0 and 1
    
    # Adjust layout to prevent legend cutoff
    plt.tight_layout()
    
    # Save the plot
    output_file = "/home/shaonian/SED/SED/script_tools/data/confidence_comparison.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\nPlot saved to: {output_file}")
    
    # Show the plot
    plt.show()

def plot_detailed_comparison():
    """Create a more detailed comparison with subplots"""
    
    data_dir = "/home/shaonian/SED/SED/script_tools/data"
    
    # Organize files by type
    teacher_files = []
    student_files = []
    
    for file in os.listdir(data_dir):
        if file.endswith("conficence.csv"):
            full_path = os.path.join(data_dir, file)
            if "teacher" in file:
                teacher_files.append(full_path)
            elif "student" in file:
                student_files.append(full_path)
    
    # Create subplots
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 12))
    
    # Plot teacher models
    colors = ['blue', 'red', 'green', 'orange']
    for i, csv_file in enumerate(teacher_files):
        try:
            df = pd.read_csv(csv_file)
            filename = os.path.basename(csv_file)
            
            if "in-domain" in filename:
                label = "Teacher - In-domain"
            elif "refTrainWeak" in filename:
                label = "Teacher - RefTrainWeak"
            else:
                label = "Teacher - Standard"
            
            ax1.plot(df['Step'], df['Value'], 
                    label=label, 
                    color=colors[i % len(colors)], 
                    linewidth=2)
            
        except Exception as e:
            print(f"Error reading {csv_file}: {e}")
    
    # Plot student models
    for i, csv_file in enumerate(student_files):
        try:
            df = pd.read_csv(csv_file)
            filename = os.path.basename(csv_file)
            
            if "in-domain" in filename:
                label = "Student - In-domain"
            elif "refTrainWeak" in filename:
                label = "Student - RefTrainWeak"
            else:
                label = "Student - Standard"
            
            ax2.plot(df['Step'], df['Value'], 
                    label=label, 
                    color=colors[i % len(colors)], 
                    linewidth=2)
            
        except Exception as e:
            print(f"Error reading {csv_file}: {e}")
    
    # Customize subplots
    ax1.set_title('Teacher Model Confidence Evolution', fontsize=14, fontweight='bold')
    ax1.set_xlabel('Training Step')
    ax1.set_ylabel('Confidence Value')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    ax1.set_ylim(0, 1.0)
    
    ax2.set_title('Student Model Confidence Evolution', fontsize=14, fontweight='bold')
    ax2.set_xlabel('Training Step')
    ax2.set_ylabel('Confidence Value')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    ax2.set_ylim(0, 1.0)
    
    plt.tight_layout()
    
    # Save the detailed plot
    output_file = "/home/shaonian/SED/SED/script_tools/data/confidence_detailed_comparison.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"Detailed plot saved to: {output_file}")
    
    plt.show()

def print_statistics():
    """Print basic statistics for each file"""
    
    data_dir = "/home/shaonian/SED/SED/script_tools/data"
    
    print("\n" + "="*80)
    print("CONFIDENCE VALUE STATISTICS")
    print("="*80)
    
    for file in os.listdir(data_dir):
        if file.endswith("conficence.csv"):
            csv_file = os.path.join(data_dir, file)
            try:
                df = pd.read_csv(csv_file)
                
                print(f"\nFile: {file}")
                print(f"  Data points: {len(df)}")
                print(f"  Step range: {df['Step'].min()} - {df['Step'].max()}")
                print(f"  Value range: {df['Value'].min():.4f} - {df['Value'].max():.4f}")
                print(f"  Mean value: {df['Value'].mean():.4f}")
                print(f"  Std deviation: {df['Value'].std():.4f}")
                print(f"  Final value: {df['Value'].iloc[-1]:.4f}")
                
            except Exception as e:
                print(f"Error processing {file}: {e}")

if __name__ == "__main__":
    print("Starting confidence value analysis and plotting...")
    
    # Print file statistics
    print_statistics()
    
    # Create plots
    print("\nCreating comparison plot...")
    plot_confidence_values()
    
    print("\nCreating detailed comparison plot...")
    plot_detailed_comparison()
    
    print("\nAnalysis complete!")
